package main

import (
	"archive/zip"
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"image"
	_ "image/png" // 导入png格式支持
	"io"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

func main() {
	walkFiles("D:\\midjourney-genernal1-4\\demo02")
}

var token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJhdWQiOiJ4ZC1kZXNpZ24tc2VydmljZSIsImV4cCI6MTcxMTEwODMyNCwiaWF0IjoxNzEwODQ5MTI0LCJpc3MiOiJ4ZC1taWRqb3VybmV5LXNlcnZlciIsIklEIjozMSwiVXNlcm5hbWUiOiIiLCJSb2xlSWRzIjpudWxsLCJQbGF0Rm9ybSI6MiwiTWVtYmVyU3RhdHVzIjoxfQ.YkDrjyFx_KWVOLhKFPjZZCij_-JzFzKZCMw5MQrZJCb9OBsLtG36x_6qJqBu9F3biDXC88YyPshGwzRJOeBxXakKn4NCcubE5txDxmx4s0zHhUfwlm7NlXGpF3-X11daZBhK7MP8cElUxT57VRpnNas6XAskIbQD7Rxzx3W6lvo"
var fileToken = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJhdWQiOiJ4ZC1kZXNpZ24iLCJleHAiOjE3MTE0NTM5MjQsImlhdCI6MTcxMDg0OTEyNCwiaXNzIjoieGQtbWlkam91cm5leS1zZXJ2ZXIiLCJUb2tlblBvbGljeSI6eyIxIjpbeyJQZXJtaXNzaW9uIjoxLCJGaWxlRGlyIjoiL2FkZGl0aW9uYWwvMzEvb3JpZ2luYWwvIn1dLCIyIjpbeyJQZXJtaXNzaW9uIjoxLCJGaWxlRGlyIjoiL3B1YmxpYy8qLyovIn0seyJQZXJtaXNzaW9uIjoxLCJGaWxlRGlyIjoiLyovMzEvKi8ifSx7IlBlcm1pc3Npb24iOjEsIkZpbGVEaXIiOiIvKi8qL3dlYl9sb3NzbGVzcy8ifSx7IlBlcm1pc3Npb24iOjEsIkZpbGVEaXIiOiIvKi8qL3dlYl9oaWdoLyJ9LHsiUGVybWlzc2lvbiI6MSwiRmlsZURpciI6Ii8qLyovd2ViX21pZGRsZS8ifSx7IlBlcm1pc3Npb24iOjEsIkZpbGVEaXIiOiIvKi8qL3dlYl9sb3cvIn1dfX0.BZaVAUy9VFl3L4B74AlyEc1GXMskAJO0GZUTsnmi1Q8z_-NmnlrN6kGqqcbPmxNObQhgShHIObG6yZ-HVtwbQuw0vScV_FiQodrP8w5DTKP-_KwkLBmH4CsUOXK4mG4kqyWrZ6mDB_9qokJB9GLirRe1dQrSLx3kTHRZTRIeJuH8on0sdxbsi9yjTL13g4J-NpvoFv4LqjQmGWYoHS7DZ6l7mQFfCVkZFyRg2W5Iqf87rowSV87HhaXgk669mhtDHy5600x6-ATEGUatp4ntx3ZnVs6I6xbM4aPoJ-MES-H8eUYteBsAAzqaTDu_UMhuAOoasdeogRk3okZp7gfmuw"

func walkFiles(folderPath string) {
	// 要遍历的文件夹路径
	//folderPath := "zip"
	//folderPath := "C:\\Users\\Administrator\\Desktop\\test-zip"
	//folderPath := "D:\\midjourney-genernal1-4\\demo01"
	// 遍历文件夹
	uploadNum := 1
	err := filepath.Walk(folderPath, func(path string, info os.FileInfo, err error) error {
		// 忽略文件夹本身和非 ZIP 文件
		if info.IsDir() || !strings.HasSuffix(strings.ToLower(path), ".zip") {
			return nil
		}

		// 解压缩 ZIP 文件
		zipReader, err := zip.OpenReader(path)
		if err != nil {
			log.Printf("failed to open ZIP file %s: %v", path, err)
			return nil
		}
		defer zipReader.Close()

		// 遍历 ZIP 文件中的文件
		for _, file := range zipReader.File {
			// 处理 JSON 文件
			if strings.HasSuffix(strings.ToLower(file.Name), ".json") {
				jsonFile, err := file.Open()
				if err != nil {
					log.Printf("failed to open JSON file %s in ZIP %s: %v", file.Name, path, err)
					continue
				}
				defer jsonFile.Close()

				// 读取并解析 JSON 文件
				var jsonData map[string]string
				if err := json.NewDecoder(jsonFile).Decode(&jsonData); err != nil {
					log.Printf("failed to decode JSON file %s in ZIP %s: %v", file.Name, path, err)
					continue
				}

				// 获取 JSON 文件名（去掉后缀）
				jsonFileName := strings.TrimSuffix(file.Name, filepath.Ext(file.Name))
				// 如果 JSON 文件名以 ".meta" 结尾，则去掉 ".meta" 后缀
				jsonFileName = strings.TrimSuffix(jsonFileName, ".meta")
				// 获取 JSON 对应的 PNG 文件名
				pngFileName := fmt.Sprintf("%s.png", jsonFileName)
				// 找到对应的 PNG 文件
				for _, zipFile := range zipReader.File {
					if zipFile.Name == pngFileName {
						// 打开 PNG 文件
						pngFile, err := zipFile.Open()
						// 使用 strings.Split 函数将字符串按照逗号分割成多个部分

						fileName := zipFile.Name
						oldSize := zipFile.FileInfo().Size()
						// 将 int64 转换为字符串
						size := strconv.FormatInt(oldSize, 10)
						url := httpUploadFile(fileName, pngFile)
						// 获取第一个部分并去掉末尾的空格

						pngFile2, _ := zipFile.Open()
						imgConfig, _, err := image.DecodeConfig(pngFile2)
						if err != nil {
							log.Fatalf("无法解码PNG文件：%v", err)
						}
						//newUrl := uploadFile(fileName, size, pngFile)
						prompt := jsonData["prompts"]
						originalUrl := fmt.Sprintf("http://10.10.22.172:9092/file/%s?token=%s&file_quality=original", url.Resource, fileToken)
						body := fmt.Sprintf(`
									{
						  "image_type": 1,
						  "member_id": 31,
						  "prompt": "%s",
						  "url": "%s",
						  "original_url": "%s",
						  "size": "%s",
						  "content_type": "image/png",
						  "width": %d,
						  "height": %d,
						  "file_name": "%s",
							 "ai_model": "Midjourney",
							"model_version":"6.0",
						}
							`, prompt, url.Resource, originalUrl, size, imgConfig.Width, imgConfig.Height, fileName)
						//fmt.Println(body)
						sendMessage(body)
						fmt.Printf("%s: upload: %v\n", folderPath, uploadNum)
						uploadNum++
						if err != nil {
							log.Printf("failed to open PNG file %s in ZIP %s: %v", pngFileName, path, err)
							continue
						}
						defer pngFile.Close()
					}
				}
			}
		}

		return nil
	})

	if err != nil {
		log.Fatal(err)
	}
}

type UploadFileResponse struct {
	Resource string `json:"resource"`
}

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("%s: %s", msg, err)
	}
}

// 连接 RabbitMQ 服务器
var conn, _ = amqp.Dial("amqp://admin:admin@10.10.22.172:5672/my_vhost")
var ch, _ = conn.Channel()

//failOnError(err, "Failed to connect to RabbitMQ")

func sendMessage(body string) {

	//defer conn.Close()

	// 创建一个 channel
	//ch, err := conn.Channel()
	//failOnError(err, "Failed to open a channel")
	//defer ch.Close()

	// 声明一个队列
	//_, err := ch.QueueDeclare(
	//	"aigcbus.gallery", // 队列名称
	//	false,             // 持久性
	//	false,             // 是否自动删除
	//	false,             // 排他性
	//	false,             // 不阻塞
	//	nil,               // 额外的属性
	//)
	//failOnError(err, "Failed to declare a queue")

	// 发送一条消息到队列
	//	body := `
	//		{
	//    "image_type": 1,
	//    "member_id": 24,
	//    "prompt": "test prompt",
	//    "url": "http://192.168.30.251:8088/testbucket/369gator369_An_idyllic_barndimonium_house_nestled_in_a_vineyard_8187cb1f-1658-41f2-840e-2f161b6710a1.png",
	//    "original_url": "http://192.168.30.251:8088/testbucket/369gator369_An_idyllic_barndimonium_house_nestled_in_a_vineyard_8187cb1f-1658-41f2-840e-2f161b6710a1.png",
	//    "size": "166666",
	//    "content_type": "image/png",
	//    "width": 1000,
	//    "height": 1000,
	//    "file_name": "aa.png",
	//	 "ai_model": "Midjourney",
	//	"model_version":"6.0",
	//}
	//`
	err := ch.Publish(
		"aigcbus.gallery", // 交换器名称
		"gallery",         // 路由键
		false,             // 强制性
		false,             // 立即发布
		amqp.Publishing{
			ContentType: "text/plain",
			Body:        []byte(body),
		})
	failOnError(err, "Failed to publish a message")

	//fmt.Println(" [x] Sent ", body)
}

// http 上传文件到AI服务器
func httpUploadFile(fileName string, file io.ReadCloser) *UploadFileResponse {

	url := "http://10.10.22.172:9092/file"
	method := "POST"
	payload := &bytes.Buffer{}
	writer := multipart.NewWriter(payload)
	_ = writer.WriteField("file_type", "additional")
	part2,
		errFile2 := writer.CreateFormFile("file", fileName)
	_, errFile2 = io.Copy(part2, file)
	if errFile2 != nil {
		fmt.Println(errFile2)
		return nil
	}
	_ = writer.WriteField("token", fileToken)
	err := writer.Close()
	if err != nil {
		fmt.Println(err)
		return nil
	}

	client := &http.Client{}
	req, err := http.NewRequest(method, url, payload)

	if err != nil {
		fmt.Println(err)
		return nil
	}
	req.Header.Add("Content-Type", "image/png")
	req.Header.Add("Authorization", token)
	req.Header.Set("Content-Type", writer.FormDataContentType())
	res, err := client.Do(req)
	if err != nil {
		fmt.Println(err)
		return nil
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		fmt.Println(err)
		return nil
	}
	//fmt.Println(string(body))

	var response = &UploadFileResponse{}
	if err := json.Unmarshal(body, response); err != nil {
		fmt.Errorf("err=%v", err)
		return nil
	}
	return response
}
